import sys
import socket
import logging
import threading
import sys, socket, time, threading
import struct
import selectors
sel = selectors.DefaultSelector()

# 接收数据缓存大小
PKT_BUFF_SIZE = 10240

logger = logging.getLogger("Proxy Logging")
formatter = logging.Formatter('%(name)-12s %(asctime)s %(levelname)-8s %(lineno)-4d %(message)s', '%Y %b %d %a %H:%M:%S',)

stream_handler = logging.StreamHandler(sys.stderr)
stream_handler.setFormatter(formatter)
logger.addHandler(stream_handler)
logger.setLevel(logging.DEBUG)

# 单向流数据传递
def tcp_mapping_worker(conn_receiver, conn_sender):
    logger.info("start ")
    while True:
        try:
            data = conn_receiver.recv(PKT_BUFF_SIZE)
        except Exception as e:
            logger.debug(e)
            logger.debug('Connection closed.')
            break

        if not data:
            logger.info('No more data is received.') # 正常情况下不会到这，一般是某个连接断了
            break

        try:
            conn_sender.sendall(data)
        except Exception:
            logger.error('Failed sending data.')
            break

        # logger.info('Mapping data > %s ' % repr(data))
        # logger.info('Mapping > %s -> %s > %d bytes.' % (conn_receiver.getpeername(), conn_sender.getpeername(), len(data)))

    conn_receiver.close()
    conn_sender.close()

    return


# 端口映射请求处理
def join(connA, connB):
    threading.Thread(target=tcp_mapping_worker, args=(connA, connB)).start()
    threading.Thread(target=tcp_mapping_worker, args=(connB, connA)).start()
    return

class frpc():
    def __init__(self, serverhost,serverport, targethost,targetport):
        self.targethost=targethost
        self.targetport=targetport
        self.serverhost=serverhost
        self.serverport=serverport
        self.server_fd = socket.create_connection((self.serverhost,self.serverport))
        self.server_fd.sendall(struct.pack('i', 1))  # 启动时 首次发送心跳包 1
        self.server_fd.setblocking(False)
        sel.register(self.server_fd,selectors.EVENT_READ,self.handle_controller_data)

        self.workConnPool = []
        threading.Thread(target=self.maitainConPool).start()
        threading.Thread(target=self.heartbeat).start()

    def heartbeat(self):
        while True:
            if self.server_fd is not None:
                self.server_fd.send(struct.pack('i', 1))
            time.sleep(9)

    # 提前把 workConn和targetConn连接起来，需要的时候直接用workConn targetConn不需要缓存
    def maitainConPool(self):
        print("启动tcp连接池")
        pool_size = 0
        while True:
            if len(self.workConnPool)<pool_size:
                workConn = socket.create_connection((self.serverhost,self.serverport))
                targetConn = socket.create_connection((self.targethost, self.targetport))
                join(targetConn,workConn)
                self.workConnPool.append(workConn)


    def handle_controller_data(self,server_fd, mask):
        try:
            data= server_fd.recv(4)  # data 长度
            if data:
                cmd = struct.unpack('i',data)[0]
                print('cmd:',cmd)
                if cmd ==2:  # 要求frpc建立的工作tcp
                    print('收到frps控制指令')
                    # 直接从池子里获取
                    if len(self.workConnPool)>0:
                        workConn = self.workConnPool.pop()
                    else:
                        targetConn = socket.create_connection((self.targethost, self.targetport))
                        workConn = socket.create_connection((self.serverhost,self.serverport))
                        join(targetConn,workConn)
                    workConn.sendall(struct.pack('i',2)) # 1 心跳包
                    print("建立工作tcp")
        except IOError as err:  # 非阻塞模式下调用 阻塞操作recv 如果没有数据会抛出异常
            # sel.unregister(conn)
            # conn.close()
            # print(err)
            pass

    def run(self):
        while True:
            events = sel.select()
            for key, mask in events:
                callback = key.data
                callback(key.fileobj, mask)
        print("frpc started!")
        
class frps(threading.Thread):
    def __init__(self, port, targetport):
        threading.Thread.__init__(self)
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.sock.bind(('0.0.0.0', port))
        # 设置为非阻塞模式
        self.sock.setblocking(False)
        self.sock.listen(100)
        sel.register(self.sock,selectors.EVENT_READ,self.accept_connection)

        frpc_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        frpc_sock.bind(('0.0.0.0', targetport))
        frpc_sock.setblocking(False)
        frpc_sock.listen(100)
        sel.register(frpc_sock,selectors.EVENT_READ,self.accept_frp_connection)


        self.frpc_cmd_conn =None
        self.userConns = []


    def heartbeat(self):
        while True:
            if self.frpc_cmd_conn is not None:
                self.frpc_cmd_conn.send(struct.pack('i', 1))
            time.sleep(9)

    # 收到用户tcp 先不接收 向frpc发送指令 让其建立工作tcp
    def accept_connection(self,sock, mask):
        userConn, addr = self.sock.accept()
        userConn.setblocking(True)
        self.userConns.append(userConn)
        print('收到用户请求')
        if self.frpc_cmd_conn is None:
            # print(1)
            return
            # time.sleep(0.5)
        print(2)
        try:
            self.frpc_cmd_conn.send(struct.pack('i',2)) # 2建立新的tcp
        except IOError as err:  # 非阻塞模式下调用 阻塞操作recv 如果没有数据会抛出异常
            print(err)
            pass


    def accept_frp_connection(self,sock, mask):
        frpc_conn, addr = sock.accept()
        frpc_conn.setblocking(False)
        # 注册为可读套接字
        sel.register(frpc_conn, selectors.EVENT_READ, self.handle_controller_data)

    def handle_controller_data(self,frpc_conn, mask):
        # print('frpc',frpc_conn,mask,self.userConns)
        try:
            data = frpc_conn.recv(4)  # Should be ready
            if data:
                cmd = struct.unpack('i',data)[0]
                print("cmd:",cmd)
                if cmd ==2:  # 是建立的工作tcp
                    sel.unregister(frpc_conn) # 不再监听
                    userConn = self.userConns.pop() # 从队列中选一个用户线程来处理
                    frpc_conn.setblocking(True)
                    join(userConn,frpc_conn)
                elif cmd ==1 and self.frpc_cmd_conn!=frpc_conn: # 说明是首次收到1
                    self.frpc_cmd_conn = frpc_conn
                    threading.Thread(target=self.heartbeat).start()
        except IOError as err:  # 非阻塞模式下调用 阻塞操作recv 如果没有数据会抛出异常
            pass

    def run(self):
        while True:
            events = sel.select()
            for key, mask in events:
                callback = key.data
                callback(key.fileobj, mask)